{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Static inverse free-boundary equilibrium calculations (in SPARC)\n",
    "\n",
    "---\n",
    "\n",
    "Here we will generate an equilibrium (find coil currents with the inverse solver) in a SPARC-like tokamak. \n",
    "\n",
    "The machine description comes from files located [here](https://github.com/cfs-energy/SPARCPublic).\n",
    "\n",
    "The equilbirium\\profile parameters are **completely made up** - please experiment on your own and change them to more realistic values as you please!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the machine object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/SPARC/SPARC_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/SPARC/SPARC_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/SPARC/SPARC_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/SPARC/SPARC_wall.pickle\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = plt.subplots(1, 1, figsize=(7, 15), dpi=80)\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.set_aspect('equal')\n",
    "tokamak.plot(axis=ax1,show=False)                                                          # plots the active coils and passive structures\n",
    "ax1.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)   # plots the limiter\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate an equilibrium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import equilibrium_update\n",
    "\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,      # provide tokamak object\n",
    "    Rmin=1.1, Rmax=2.7,   # radial range\n",
    "    Zmin=-1.8, Zmax=1.8,  # vertical range\n",
    "    nx=129,                # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,                # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    # psi=plasma_psi\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate a profile object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise the profiles\n",
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,        # equilibrium object\n",
    "    paxis=5e4,    # pressure on axis\n",
    "    Ip=8.7e6,       # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=1.8,  # profile function parameter\n",
    "    alpha_n=1.2   # profile function parameter\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the static nonlinear solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rx = 1.55      # X-point radius\n",
    "Zx = 1.15      # X-point height\n",
    "\n",
    "# set desired null_points locations\n",
    "# this can include X-point and O-point locations\n",
    "null_points = [[Rx, Rx], [Zx, -Zx]]\n",
    "\n",
    "Rout = 2.4    # outboard midplane radius\n",
    "Rin = 1.3    # inboard midplane radius\n",
    "\n",
    "# set desired isoflux constraints with format \n",
    "# isoflux_set = [isoflux_0, isoflux_1 ... ] \n",
    "# with each isoflux_i = [R_coords, Z_coords]\n",
    "isoflux_set = np.array([[[Rx, Rx, Rout, Rin, 1.7, 1.7], [Zx, -Zx, 0.0, 0.0, 1.5, -1.5]]])\n",
    "           \n",
    "# instantiate the freegsnke constrain object\n",
    "from freegsnke.inverse import Inverse_optimizer\n",
    "constrain = Inverse_optimizer(null_points=null_points,\n",
    "                              isoflux_set=isoflux_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The inverse solve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GSStaticSolver.inverse_solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=constrain, \n",
    "                     target_relative_tolerance=1e-6,\n",
    "                     target_relative_psit_update=1e-3,\n",
    "                     verbose=False, # print output\n",
    "                     l2_reg=np.array([1e-16]*10 + [1e-5]), \n",
    "                     )\n",
    "                     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# refine with a forward solve (now the currents are known)\n",
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9,\n",
    "                     verbose=False, # print output\n",
    "                     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = plt.subplots(1, 1, figsize=(7, 15), dpi=80)\n",
    "\n",
    "ax1.grid(zorder=0, alpha=0.75)\n",
    "ax1.set_aspect('equal')\n",
    "eq.tokamak.plot(axis=ax1,show=False)                                                          # plots the active coils and passive structures\n",
    "ax1.fill(tokamak.wall.R, tokamak.wall.Z, color='k', linewidth=1.2, facecolor='w', zorder=0)   # plots the limiter\n",
    "eq.plot(axis=ax1,show=False)                                                                  # plots the equilibrium\n",
    "constrain.plot(axis=ax1, show=False)                                                          # plots the contraints\n",
    "ax1.set_xlim(1.0, 3.0)\n",
    "ax1.set_ylim(-2.0, 2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eq.tokamak.getCurrents()\n",
    "\n",
    "# # save coil currents to file\n",
    "# import pickle\n",
    "# with open('simple_diverted_currents_PaxisIp.pk', 'wb') as f:\n",
    "#     pickle.dump(obj=inverse_current_values, file=f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke4e",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}